package com.study.springmvc.config;

import org.springframework.lang.NonNull;

import java.util.concurrent.*;

public class CustomThreadPoolExecutor<T> extends ThreadPoolExecutor {

    private final ThreadLocalContext<T> threadLocalContext;

    public CustomThreadPoolExecutor(ThreadLocalContext<T> context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        this.threadLocalContext = context;
    }

    public CustomThreadPoolExecutor(ThreadLocalContext<T> context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
        this.threadLocalContext = context;
    }

    public CustomThreadPoolExecutor(ThreadLocalContext<T> context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
        this.threadLocalContext = context;
    }

    public CustomThreadPoolExecutor(ThreadLocalContext<T> context, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        this.threadLocalContext = context;
    }

    @Override
    public void execute(@NonNull Runnable command) {
        super.execute(new RunnableWrapper(command));
    }

    public interface ThreadLocalContext<T> {
        T getContext();

        void setContext(T context);
    }

    private class RunnableWrapper implements Runnable {
        private final Runnable target;
        private final T context;

        public RunnableWrapper(Runnable target) {
            this.target = target;
            this.context = threadLocalContext.getContext();
        }

        @Override
        public void run() {
            CustomThreadPoolExecutor.this.threadLocalContext.setContext(context);
            target.run();
        }
    }
}
